"""经典技术指标计算函数."""

from __future__ import annotations

import numpy as np
import pandas as pd

from .utils import _ema, _rolling_std, _sma


def _macd(series: pd.Series, fast: int = 12, slow: int = 26, signal: int = 9) -> pd.DataFrame:
    """计算 MACD 指标（Moving Average Convergence Divergence）.
    
    MACD是趋势跟踪动量指标，通过快慢EMA的差值来判断趋势变化。
    
    Parameters
    ----------
    series:
        价格序列
    fast:
        快线周期（默认12）
    slow:
        慢线周期（默认26）
    signal:
        信号线周期（默认9）
        
    Returns
    -------
    pd.DataFrame
        包含 macd_dif、macd_dea、macd_hist 的DataFrame
    """
    ema_fast = _ema(series, fast)
    ema_slow = _ema(series, slow)
    dif = ema_fast - ema_slow
    dea = _ema(dif, signal)
    hist = dif - dea
    return pd.DataFrame({"macd_dif": dif, "macd_dea": dea, "macd_hist": hist})


def _rsi(series: pd.Series, window: int = 14) -> pd.Series:
    """计算 RSI 指标（Relative Strength Index）.
    
    RSI是动量振荡器，用于衡量价格变动的速度和幅度，范围0-100。
    
    Parameters
    ----------
    series:
        价格序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.Series
        RSI序列（0-100）
    """
    delta = series.diff()
    gain = delta.clip(lower=0).rolling(window, min_periods=window).mean()
    loss = -delta.clip(upper=0).rolling(window, min_periods=window).mean()
    rs = gain / loss.replace(0, np.nan)
    return 100 - (100 / (1 + rs))


def _cmo(series: pd.Series, window: int = 14) -> pd.Series:
    """计算 CMO（Chande Momentum Oscillator）.
    
    CMO是动量振荡器，与RSI类似但计算方法不同，范围-100到+100。
    
    Parameters
    ----------
    series:
        价格序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.Series
        CMO序列（-100到+100）
    """
    delta = series.diff()

    gains = delta.clip(lower=0).rolling(window, min_periods=window).sum()
    losses = (-delta.clip(upper=0)).rolling(window, min_periods=window).sum()

    denominator = gains + losses
    cmo = 100 * (gains - losses) / denominator.replace(0, np.nan)
    return cmo


def _atr(high: pd.Series, low: pd.Series, close: pd.Series, window: int = 14) -> pd.Series:
    """计算 ATR（Average True Range）.
    
    ATR衡量市场波动性，基于真实波动范围（TR）的平均值。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.Series
        ATR序列
    """
    prev_close = close.shift(1)
    tr = pd.concat(
        [
            high - low,
            (high - prev_close).abs(),
            (low - prev_close).abs(),
        ],
        axis=1,
    ).max(axis=1)
    return tr.rolling(window, min_periods=window).mean()


def _bollinger(series: pd.Series, window: int = 20, num_std: float = 2.0) -> pd.DataFrame:
    """计算布林带（Bollinger Bands）.
    
    布林带由中轨（SMA）、上轨（中轨+标准差）、下轨（中轨-标准差）组成，
    用于判断价格是否偏离正常范围。
    
    Parameters
    ----------
    series:
        价格序列
    window:
        窗口大小（默认20）
    num_std:
        标准差倍数（默认2.0）
        
    Returns
    -------
    pd.DataFrame
        包含 bb_mid、bb_upper、bb_lower、bb_width、bb_position 的DataFrame
    """
    mid = _sma(series, window)
    std = _rolling_std(series, window)
    upper = mid + num_std * std
    lower = mid - num_std * std
    width = (upper - lower) / mid
    pos = (series - lower) / (upper - lower)
    return pd.DataFrame(
        {
            "bb_mid": mid,
            "bb_upper": upper,
            "bb_lower": lower,
            "bb_width": width,
            "bb_position": pos,
        }
    )


def _kdj(high: pd.Series, low: pd.Series, close: pd.Series, n: int = 9, m1: int = 3, m2: int = 3) -> pd.DataFrame:
    """计算 KDJ 指标（随机指标）.
    
    KDJ是随机振荡器，用于识别超买超卖区域，由K、D、J三条线组成。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    n:
        RSV计算周期（默认9）
    m1:
        K值平滑周期（默认3）
    m2:
        D值平滑周期（默认3）
        
    Returns
    -------
    pd.DataFrame
        包含 kdj_k、kdj_d、kdj_j 的DataFrame
    """
    rsv = (close - low.rolling(n, min_periods=n).min()) / (
        high.rolling(n, min_periods=n).max() - low.rolling(n, min_periods=n).min()
    ) * 100
    rsv = rsv.replace([np.inf, -np.inf], np.nan)
    k = rsv.ewm(alpha=1 / m1, adjust=False).mean()
    d = k.ewm(alpha=1 / m2, adjust=False).mean()
    j = 3 * k - 2 * d
    return pd.DataFrame({"kdj_k": k, "kdj_d": d, "kdj_j": j})


def _cci(high: pd.Series, low: pd.Series, close: pd.Series, window: int = 20) -> pd.Series:
    """计算 CCI（商品通道指数）.
    
    CCI衡量价格偏离统计平均值的程度，用于识别超买超卖。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    window:
        窗口大小（默认20）
        
    Returns
    -------
    pd.Series
        CCI序列
    """
    tp = (high + low + close) / 3
    sma_tp = _sma(tp, window)
    mad = tp.rolling(window, min_periods=window).apply(lambda x: np.abs(x - x.mean()).mean(), raw=False)
    cci = (tp - sma_tp) / (0.015 * mad.replace(0, np.nan))
    return cci.replace([np.inf, -np.inf], np.nan)


def _adx(high: pd.Series, low: pd.Series, close: pd.Series, window: int = 14) -> pd.DataFrame:
    """计算 ADX（平均趋向指数）及相关指标 +DI、-DI.
    
    ADX衡量趋势强度，+DI和-DI分别表示上升和下降趋势的强度。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.DataFrame
        包含 adx、adx_plus_di、adx_minus_di 的DataFrame
    """
    high_diff = high.diff()
    low_diff = -low.diff()
    
    plus_dm = high_diff.where((high_diff > low_diff) & (high_diff > 0), 0)
    minus_dm = low_diff.where((low_diff > high_diff) & (low_diff > 0), 0)
    
    atr_val = _atr(high, low, close, window)
    
    plus_di = 100 * (plus_dm.rolling(window, min_periods=window).mean() / atr_val.replace(0, np.nan))
    minus_di = 100 * (minus_dm.rolling(window, min_periods=window).mean() / atr_val.replace(0, np.nan))
    
    dx = 100 * np.abs(plus_di - minus_di) / (plus_di + minus_di).replace(0, np.nan)
    adx = dx.rolling(window, min_periods=window).mean()
    
    return pd.DataFrame({"adx": adx, "adx_plus_di": plus_di, "adx_minus_di": minus_di})


def _obv(close: pd.Series, volume: pd.Series) -> pd.Series:
    """计算 OBV（能量潮指标）.
    
    OBV通过成交量变化预测价格变化，基于价涨量增、价跌量减的假设。
    
    Parameters
    ----------
    close:
        收盘价序列
    volume:
        成交量序列
        
    Returns
    -------
    pd.Series
        OBV序列
    """
    price_change = close.diff()
    obv = (volume * np.sign(price_change)).fillna(0).cumsum()
    return obv


def _mfi(high: pd.Series, low: pd.Series, close: pd.Series, volume: pd.Series, window: int = 14) -> pd.Series:
    """计算 MFI（资金流量指标）.
    
    MFI是结合价格和成交量的动量指标，类似于RSI但考虑了成交量因素。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    volume:
        成交量序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.Series
        MFI序列（0-100）
    """
    typical_price = (high + low + close) / 3
    raw_money_flow = typical_price * volume
    
    positive_flow = raw_money_flow.where(typical_price > typical_price.shift(1), 0)
    negative_flow = raw_money_flow.where(typical_price < typical_price.shift(1), 0)
    
    positive_sum = positive_flow.rolling(window, min_periods=window).sum()
    negative_sum = negative_flow.rolling(window, min_periods=window).sum()
    
    money_flow_ratio = positive_sum / negative_sum.replace(0, np.nan)
    mfi = 100 - (100 / (1 + money_flow_ratio))
    
    return mfi.replace([np.inf, -np.inf], np.nan)


def _williams_r(high: pd.Series, low: pd.Series, close: pd.Series, window: int = 14) -> pd.Series:
    """计算 Williams %R 指标.
    
    Williams %R是动量振荡器，用于识别超买超卖区域，范围-100到0。
    
    Parameters
    ----------
    high:
        最高价序列
    low:
        最低价序列
    close:
        收盘价序列
    window:
        窗口大小（默认14）
        
    Returns
    -------
    pd.Series
        Williams %R序列（-100到0）
    """
    highest_high = high.rolling(window, min_periods=window).max()
    lowest_low = low.rolling(window, min_periods=window).min()
    wr = -100 * (highest_high - close) / (highest_high - lowest_low).replace(0, np.nan)
    return wr.replace([np.inf, -np.inf], np.nan)

